#pragma once
#include <rclcpp/rclcpp.hpp>
#include <nav_msgs/msg/odometry.hpp>
#include <tf2_ros/transform_broadcaster.h>
#include <navigation/data_struct/pose_info.hpp>
#include <tf2/LinearMath/Quaternion.h>
#include <visualization_msgs/msg/marker.hpp>
#include <visualization_msgs/msg/marker_array.hpp>
#include <navigation/data_struct/linked_list.hpp>
#include <navigation/msg/targets_list.hpp>
#include <navigation/nav_core.hpp>
#include <std_msgs/msg/string.hpp>
#include <filesystem>
#include <fstream>

class StateManager {
public:
    explicit StateManager(
        const rclcpp::Node::SharedPtr& node,
        std::string maps_dir_path,
        double visualization_hz = 10.0
    ):
    node_(node), maps_dir_path_(maps_dir_path),
    visualization_hz_(visualization_hz) {
        adjusted_map_pose_ = std::make_shared<PoseInfo>("camera_init", "adjusted_init");
        adjusted_map_sub_ = node->create_subscription<geometry_msgs::msg::Pose>(
            "/adjusted_map_pose", 10,
            std::bind(&StateManager::adjusted_map_callback, this, std::placeholders::_1)
        );  // 获取scripts/map_adjuster.py通过LightGlue调整后的地图位姿
        odom_sub_ = node_->create_subscription<nav_msgs::msg::Odometry>(
            "/Odometry", 10,
            std::bind(&StateManager::odom_callback, this, std::placeholders::_1)
        );  // 获取fast lio的camera_init->body的Odom信息
        RCLCPP_INFO(node_->get_logger(), "Start publishing robot_init, robot TF and robot_marker.");
        robot_marker_pub_ = node_->create_publisher<visualization_msgs::msg::Marker>("/robot_marker", 1);
        target_marker_pub_ = node_->create_publisher<visualization_msgs::msg::MarkerArray>("/target_markers_array", 1);
        robot_tf_broadcaster_ = std::make_shared<tf2_ros::TransformBroadcaster>(node_);

        robot2targets.set_head(PoseInfo("robot_init", "robot", 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0));
        target_names_pub_ = node_->create_publisher<navigation::msg::TargetsList>("/target_names", 1);
        target_thread_ = std::thread(&StateManager::target_thread_func, this);
    }
    ~StateManager() {
        target_thread_running_ = false;
        if (target_thread_.joinable()) target_thread_.join();
    }
    void add_target(
        double x, double y, double z,
        double qx, double qy, double qz, double qw,
        int index = 0  // 默认添加为当前第一个目标点
    ) {
        auto target_name = std::string("target_") + std::to_string(target_count_++);
        auto target = PoseInfo("robot_init", target_name, x, y, z, qx, qy, qz, qw);
        robot2targets.add_node(target, index);
        RCLCPP_INFO(node_->get_logger(), "Adding %dth target at index %d: (%.2f, %.2f, %.2f), (%.2f, %.2f, %.2f, %.2f)",
            target_count_, index, x, y, z, qx, qy, qz, qw);
    }
    void remove_target(int index) {
        bool result = robot2targets.pop_node(index);
        RCLCPP_INFO(node_->get_logger(), "Removing target at index %d: %s", index, result ? "Success" : "Failed");
    }
    void move_target(int index1, int index2) {
        bool result = robot2targets.move_node(index1, index2);
        RCLCPP_INFO(node_->get_logger(), "Moving target from index %d to %d: %s", index1, index2, result ? "Success" : "Failed");
    }
    void set_highlight_targets(std::vector<int> &indices) {
        std::string info;
        highlighted_targets_.clear();
        for (auto index: indices) {
            highlighted_targets_.insert(index);
            info += std::to_string(index) + " ";
        }
        RCLCPP_INFO(node_->get_logger(), "Highlighted targets: %s", info.c_str());
    }
    std::pair<bool, std::string> save_targets_to_file(std::string file_name) {
        std::string message;
        std::string file_path = std::filesystem::path(maps_dir_path_) / file_name;
        std::ofstream file(file_path);
        if (!file.is_open()) {
            message = "Failed to open targets file: " + file_path;
            RCLCPP_ERROR(node_->get_logger(), message.c_str());
            return std::make_pair(false, message);
        }
        auto node = robot2targets.second_node();
        while (node != nullptr) {
            auto& pose = node->data;
            file << pose.point.x << " " << pose.point.y << " " << pose.point.z << " "
                 << pose.quaternion.x << " " << pose.quaternion.y << " "
                 << pose.quaternion.z << " " << pose.quaternion.w << "\n";
            node = node->next;
        }
        message = "Saved targets to file: " + file_path;
        RCLCPP_INFO(node_->get_logger(), message.c_str());
        file.close();
        return std::make_pair(true, message);
    }
    std::pair<bool, std::string> load_targets_from_file(std::string file_path) {
        std::string message;
        std::ifstream file(file_path);
        if (!file.is_open()) {
            message = "Failed to open targets file: " + file_path;
            RCLCPP_INFO(node_->get_logger(), message.c_str());
            return std::make_pair(false, message);
        }
        robot2targets.clear(1);  // 清空targets
        target_count_ = 0;
        std::string line;
        while (std::getline(file, line)) {
            std::istringstream iss(line);
            double x, y, z, qx, qy, qz, qw;
            if (!(iss >> x >> y >> z >> qx >> qy >> qz >> qw)) {
                RCLCPP_ERROR(node_->get_logger(), "Invalid target line: %s", line.c_str());
                continue;
            }
            add_target(x, y, z, qx, qy, qz, qw, -1);
        }
        message = "Loaded targets from file (" + std::to_string(target_count_) + "targets): " + file_path;
        RCLCPP_INFO(node_->get_logger(), message.c_str());
        file.close();
        return std::make_pair(true, message);
    }
    std::string get_maps_dir_path() { return maps_dir_path_; }
    LinkedList<PoseInfo> get_robot2targets() {
        return robot2targets;
    }
    void change_target_threashold(double dist_threshold, double quat_threshold) {
        target_dist_threshold_ = dist_threshold;
        target_yaw_threshold_ = quat_threshold;
        RCLCPP_INFO(node_->get_logger(), "Changed target thresholds to distance: %.2f, angle: %.2f",
            target_dist_threshold_.load(), target_yaw_threshold_.load());
    }
    std::pair<double, double> get_target_thresholds() const {
        return {target_dist_threshold_.load(), target_yaw_threshold_.load()};
    }

private:
    /**
     * @brief 创建一个可视化Marker消息，用于在RViz中显示位姿等信息。
     * 
     * @param pose  位置和方向信息（PoseInfo类型）
     * @param frame_id  坐标系ID
     * @param ns  Marker命名空间
     * @param id  Marker的唯一ID
     * @param scale  Marker的尺寸（长度, 宽度, 高度），默认{0.5, 0.1, 0.1}
     * @param color  Marker的颜色（R, G, B），默认{1.0, 0.0, 0.0}
     * @return visualization_msgs::msg::Marker  创建的Marker消息
     */
    visualization_msgs::msg::Marker create_marker(
        PoseInfo pose, std::string frame_id,
        std::string ns, int id,
        int marker_type = visualization_msgs::msg::Marker::ARROW,
        std::tuple<double, double, double> scale = {0.5, 0.1, 0.1},
        std::tuple<double, double, double> color = {1.0, 0.0, 0.0},
        rclcpp::Time stamp = rclcpp::Time(0, 0),
        rclcpp::Duration lifetime = rclcpp::Duration::from_seconds(0.1)
    ) {
        visualization_msgs::msg::Marker marker;
        marker.header.stamp = stamp;  // 使用frame_id TF最新发布的时间
        marker.header.frame_id = frame_id;
        marker.ns = ns;
        marker.id = id;
        marker.type = marker_type;
        marker.action = visualization_msgs::msg::Marker::ADD;
        marker.pose = pose.to_pose_msg();
        marker.scale.x = std::get<0>(scale);  // 长度
        marker.scale.y = std::get<1>(scale);  // 宽度
        marker.scale.z = std::get<2>(scale);  // 高度
        marker.color.r = std::get<0>(color);
        marker.color.g = std::get<1>(color);
        marker.color.b = std::get<2>(color);
        marker.color.a = 1.0;  // 不透明
        marker.lifetime = lifetime;
        return marker;
    }

    /**
     * @brief 处理从/Odometry Topic接收到的消息，发布robot_init, robot的TF和robot_marker
     * @brief 并不断更新robot_2_targets_中第一个节点, 即机器人当前位置
     * @param msg 接收到的/Odometry消息指针
     */
    void odom_callback(const nav_msgs::msg::Odometry::SharedPtr msg) {
        auto& pos = msg->pose.pose.position;
        auto& orient = msg->pose.pose.orientation;
        auto& adjusted = adjusted_map_pose_;
        auto t_adjusted = adjusted->to_tf_transform();

        // camera_init -> adjusted_init的TF
        auto tm_adjusted = adjusted->to_tf_msg();
        tm_adjusted.header.stamp = msg->header.stamp;
        robot_tf_broadcaster_->sendTransform(tm_adjusted);

        // adjusted_init -> robot_init的TF, 不移动仅绕z轴旋转180度
        tf2::Quaternion q_rot180;
        q_rot180.setRPY(0.0, 0.0, M_PI);
        auto t_rot180 = tf2::Transform(q_rot180, {0.0, 0.0, 0.0});
        auto t_robot_init = t_rot180.inverse();
        auto robot_init_pose = PoseInfo("adjusted_init", "robot_init", t_robot_init);
        auto tm_robot_init = robot_init_pose.to_tf_msg();
        tm_robot_init.header.stamp = msg->header.stamp;
        robot_tf_broadcaster_->sendTransform(tm_robot_init);

        // 要从camera_init -> body -> robot转换得到
        // 而camera_init -> adjusted -> robot_init
        // 因此robot_init -> robot变换为 R_rot180_inv * adjusted * R_body * R_rot180
        // x, y翻转即可
        tf2::Quaternion q_body = { orient.x, orient.y, orient.z, orient.w };
        double roll, pitch, yaw;
        tf2::Matrix3x3(q_body).getRPY(roll, pitch, yaw);
        tf2::Quaternion q_body_regular;
        q_body_regular.setRPY(0.0, 0.0, yaw);
        auto t_body_regular = tf2::Transform(q_body_regular, {pos.x, pos.y, pos.z});
        auto t_robot = t_rot180.inverse() * t_adjusted.inverse() * t_body_regular * t_rot180;
        auto robot_pose = PoseInfo("robot_init", "robot", t_robot);
        auto tm_robot = robot_pose.to_tf_msg();
        tm_robot.header.stamp = msg->header.stamp;
        robot_tf_broadcaster_->sendTransform(tm_robot);

        // 发布robot的marker
        auto robot_marker = create_marker(
            robot_pose,
            "robot_init", "robot", 1,
            visualization_msgs::msg::Marker::ARROW,
            {0.5, 0.1, 0.1}, {0.0, 1.0, 0.0},  // 绿色箭头
            msg->header.stamp
        );
        robot_marker_pub_->publish(robot_marker);

        // 更新robot_2_targets_的第一个节点 (当前机器人位置)
        robot2targets.head()->data = robot_pose;
        
        last_odom_time_ = msg->header.stamp;
    }

    void adjusted_map_callback(const geometry_msgs::msg::Pose::SharedPtr msg) {
        auto pos = msg->position;
        auto orient = msg->orientation;
        *adjusted_map_pose_ = PoseInfo(
            "camera_init", "adjusted_init",
            pos.x, pos.y, pos.z,
            orient.x, orient.y, orient.z, orient.w
        );
        RCLCPP_INFO(node_->get_logger(), "Adjusted map pose updated: x=%.2f, y=%.2f, z=%.2f, qx=%.2f, qy=%.2f, qz=%.2f, qw=%.2f",
            pos.x, pos.y, pos.z, orient.x, orient.y, orient.z, orient.w);
    }

    /**
     * @brief 可视化当前的所有导航目标点, 用line_strip连接路径, 每个目标点用箭头表示,
     * @brief 并检查第一个目标点是否到达, 如果到达则删除该目标点
     */
    void target_thread_func() {
        RCLCPP_INFO(node_->get_logger(), "Start target thread for visualizing targets.");
        rclcpp::Rate rate(visualization_hz_);
        while (rclcpp::ok() && target_thread_running_) {
            // 检查第一个目标点是否到达
            auto first_target = robot2targets.second_node();
            if (first_target != nullptr) {
                auto robot_pos = robot2targets.head()->data;
                auto target_pos = first_target->data;

                auto [dist_norm, quat_norm] = (robot_pos - target_pos).norm();

                if (dist_norm < target_dist_threshold_ && quat_norm < target_yaw_threshold_) {
                    RCLCPP_INFO(node_->get_logger(), "%s reached, removing first target.", 
                        first_target->data.child_frame_id.c_str());
                    robot2targets.pop_node(1);
                } else {
                    // RCLCPP_INFO(node_->get_logger(), "%s not reached yet, distance: %.2f, angle: %.2f",
                    //     first_target->data.child_frame_id.c_str(), dist_norm, quat_norm);
                }
            }

            navigation::msg::TargetsList target_name_msg;
            if (robot2targets.size() <= 1) {
                target_names_pub_->publish(target_name_msg);
                rate.sleep();
                continue;
            }

            auto pre_node = robot2targets.head();  // robot当前位置节点
            auto now_node = robot2targets.second_node();  // 第一个目标点

            visualization_msgs::msg::MarkerArray marker_array;

            for (int i = 0; now_node != nullptr; ++i) {
                visualization_msgs::msg::Marker target_marker;
                if (highlighted_targets_.find(i+1) != highlighted_targets_.end()) {
                    target_marker = create_marker(
                        now_node->data, "robot_init", "target_highlighted", i,
                        visualization_msgs::msg::Marker::ARROW,
                        {0.6, 0.2, 0.2}, {1.0, 1.0, 0.0},  // 黄色箭头
                        last_odom_time_
                    );
                } else {
                    target_marker = create_marker(
                        now_node->data, "robot_init", "target", i,
                        visualization_msgs::msg::Marker::ARROW,
                        {0.3, 0.08, 0.08}, {1.0, 0.0, 0.0},
                        last_odom_time_
                    );
                }
                auto line_marker = create_marker(
                    now_node->data, "robot_init", "line", i,
                    visualization_msgs::msg::Marker::ARROW,
                    {0.05, 0.0, 1e-3}, {0.0, 0.0, 1.0}  // 蓝色线条
                );
                line_marker.pose = geometry_msgs::msg::Pose();  // reset to empty
                line_marker.points.push_back(pre_node->data.to_pose_msg().position);
                line_marker.points.push_back(now_node->data.to_pose_msg().position);
                marker_array.markers.push_back(target_marker);
                marker_array.markers.push_back(line_marker);
                target_name_msg.names.push_back(now_node->data.child_frame_id);

                pre_node = now_node;
                now_node = now_node->next;
            }
            target_marker_pub_->publish(marker_array);
            target_names_pub_->publish(target_name_msg);

            rate.sleep();
        }
    }

    // 初始化参数
    rclcpp::Node::SharedPtr node_;
    std::string maps_dir_path_;
    double visualization_hz_;

    // ROS相关
    rclcpp::Subscription<nav_msgs::msg::Odometry>::SharedPtr odom_sub_;
    rclcpp::Subscription<geometry_msgs::msg::Pose>::SharedPtr adjusted_map_sub_;
    rclcpp::Publisher<visualization_msgs::msg::Marker>::SharedPtr robot_marker_pub_;
    rclcpp::Publisher<visualization_msgs::msg::MarkerArray>::SharedPtr target_marker_pub_;
    rclcpp::Publisher<navigation::msg::TargetsList>::SharedPtr target_names_pub_;
    std::shared_ptr<tf2_ros::TransformBroadcaster> robot_tf_broadcaster_;
    rclcpp::Time last_odom_time_;

    // target相关
    std::atomic<bool> target_thread_running_{true};
    LinkedList<PoseInfo> robot2targets;
    std::thread target_thread_;
    int target_count_{0};
    std::atomic<double> target_dist_threshold_{0.1};
    std::atomic<double> target_yaw_threshold_{M_PI / 18.0};  // 10度
    std::set<int> highlighted_targets_;

    // adjusted map相关
    std::shared_ptr<PoseInfo> adjusted_map_pose_;
};
